package com.video.tool;

import org.bytedeco.ffmpeg.global.avcodec;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.FFmpegFrameRecorder;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;

import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.ConvolveOp;
import java.awt.image.Kernel;

public class AutoVideoFilterProcessor {

    private static final double BRIGHTNESS_THRESHOLD_LOW = 0.3;
    private static final double BRIGHTNESS_THRESHOLD_HIGH = 0.7;
    private static final double CONTRAST_THRESHOLD = 0.5;
    private static final double NOISE_THRESHOLD = 0.1;
    private static final double COLOR_VARIANCE_THRESHOLD = 0.15;

    public static void main(String[] args) {
        String inputVideo = "F:\\input.mp4";
        String outputVideo = "F:\\output.mp4";

        try {
            applyAutoFilter(inputVideo, outputVideo);
            System.out.println("finish");
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static void applyAutoFilter(String inputVideo, String outputVideo) throws Exception {
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(inputVideo);
        grabber.start();

        FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(outputVideo,
                grabber.getImageWidth(),
                grabber.getImageHeight());

        recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
        recorder.setFormat("mp4");
        recorder.setFrameRate(grabber.getFrameRate());
        recorder.setVideoBitrate(grabber.getVideoBitrate());
        recorder.start();

        Java2DFrameConverter converter = new Java2DFrameConverter();
        Frame frame;

        while ((frame = grabber.grab()) != null) {
            if (frame.image != null) {
                BufferedImage bufferedImage = converter.convert(frame);

                ImageFeatures features = analyzeImageFeatures(bufferedImage);

                BufferedImage processedImage = applyOptimalFilter(bufferedImage, features);

                Frame processedFrame = converter.convert(processedImage);
                recorder.record(processedFrame);
            }
        }

        grabber.stop();
        grabber.release();
        recorder.stop();
        recorder.release();
    }

    private static class ImageFeatures {
        double averageBrightness;
        double contrastLevel;
        double noiseLevel;
        double colorVariance;
        boolean isUnderexposed;
        boolean isOverexposed;
        boolean isLowContrast;
        boolean isNoisy;
        boolean hasHighColorVariance;
    }

    // 分析图像特征
    private static ImageFeatures analyzeImageFeatures(BufferedImage image) {
        ImageFeatures features = new ImageFeatures();
        int width = image.getWidth();
        int height = image.getHeight();
        long totalBrightness = 0;
        long totalVariance = 0;
        int[] brightnessValues = new int[width * height];
        int pixelCount = 0;

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                Color color = new Color(image.getRGB(x, y));
                int brightness = (color.getRed() + color.getGreen() + color.getBlue()) / 3;
                brightnessValues[pixelCount] = brightness;
                totalBrightness += brightness;
                pixelCount++;
            }
        }

        features.averageBrightness = (double) totalBrightness / (width * height * 255);

        double meanBrightness = (double) totalBrightness / pixelCount;
        for (int i = 0; i < pixelCount; i++) {
            totalVariance += Math.pow(brightnessValues[i] - meanBrightness, 2);
        }
        features.contrastLevel = Math.sqrt(totalVariance / pixelCount) / 255;

        double totalNoise = 0;
        for (int y = 1; y < height - 1; y++) {
            for (int x = 1; x < width - 1; x++) {
                Color center = new Color(image.getRGB(x, y));
                Color next = new Color(image.getRGB(x + 1, y));
                totalNoise += getColorDifference(center, next);
            }
        }
        features.noiseLevel = totalNoise / (width * height);

        // 计算色彩方差
        features.colorVariance = calculateColorVariance(image);

        // 设置特征标志
        features.isUnderexposed = features.averageBrightness < BRIGHTNESS_THRESHOLD_LOW;
        features.isOverexposed = features.averageBrightness > BRIGHTNESS_THRESHOLD_HIGH;
        features.isLowContrast = features.contrastLevel < CONTRAST_THRESHOLD;
        features.isNoisy = features.noiseLevel > NOISE_THRESHOLD;
        features.hasHighColorVariance = features.colorVariance > COLOR_VARIANCE_THRESHOLD;

        return features;
    }

    private static double getColorDifference(Color c1, Color c2) {
        int rDiff = c1.getRed() - c2.getRed();
        int gDiff = c1.getGreen() - c2.getGreen();
        int bDiff = c1.getBlue() - c2.getBlue();
        return Math.sqrt(rDiff * rDiff + gDiff * gDiff + bDiff * bDiff) / 441.67; // 归一化到0-1
    }

    private static double calculateColorVariance(BufferedImage image) {
        long totalR = 0, totalG = 0, totalB = 0;
        int width = image.getWidth();
        int height = image.getHeight();
        int pixels = width * height;

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                Color color = new Color(image.getRGB(x, y));
                totalR += color.getRed();
                totalG += color.getGreen();
                totalB += color.getBlue();
            }
        }

        double avgR = (double) totalR / pixels;
        double avgG = (double) totalG / pixels;
        double avgB = (double) totalB / pixels;

        double varianceR = 0, varianceG = 0, varianceB = 0;
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                Color color = new Color(image.getRGB(x, y));
                varianceR += Math.pow(color.getRed() - avgR, 2);
                varianceG += Math.pow(color.getGreen() - avgG, 2);
                varianceB += Math.pow(color.getBlue() - avgB, 2);
            }
        }

        return (Math.sqrt(varianceR + varianceG + varianceB) / pixels) / 255;
    }

    private static BufferedImage applyOptimalFilter(BufferedImage source, ImageFeatures features) {
        BufferedImage result = source;

        if (features.isUnderexposed) {
            result = applyBrightnessFilter(result, 1.3f);
        }

        if (features.isOverexposed) {
            result = applyBrightnessFilter(result, 0.8f);
        }

        if (features.isLowContrast) {
            result = applyContrastFilter(result, 1.4f);
        }

        if (features.isNoisy) {
            result = applyDenoiseFilter(result);
        }

        if (features.hasHighColorVariance) {
            result = applyColorBalanceFilter(result);
        }

        return result;
    }

    // 亮度调整滤镜
    private static BufferedImage applyBrightnessFilter(BufferedImage source, float factor) {
        BufferedImage result = new BufferedImage(
                source.getWidth(),
                source.getHeight(),
                BufferedImage.TYPE_INT_RGB);

        for (int y = 0; y < source.getHeight(); y++) {
            for (int x = 0; x < source.getWidth(); x++) {
                Color color = new Color(source.getRGB(x, y));

                int r = Math.min(255, Math.max(0, (int) (color.getRed() * factor)));
                int g = Math.min(255, Math.max(0, (int) (color.getGreen() * factor)));
                int b = Math.min(255, Math.max(0, (int) (color.getBlue() * factor)));

                result.setRGB(x, y, new Color(r, g, b).getRGB());
            }
        }
        return result;
    }

    private static BufferedImage applyContrastFilter(BufferedImage source, float factor) {
        BufferedImage result = new BufferedImage(
                source.getWidth(),
                source.getHeight(),
                BufferedImage.TYPE_INT_RGB);

        for (int y = 0; y < source.getHeight(); y++) {
            for (int x = 0; x < source.getWidth(); x++) {
                Color color = new Color(source.getRGB(x, y));

                int r = color.getRed();
                int g = color.getGreen();
                int b = color.getBlue();

                r = (int) (((r / 255.0 - 0.5) * factor + 0.5) * 255.0);
                g = (int) (((g / 255.0 - 0.5) * factor + 0.5) * 255.0);
                b = (int) (((b / 255.0 - 0.5) * factor + 0.5) * 255.0);

                r = Math.max(0, Math.min(255, r));
                g = Math.max(0, Math.min(255, g));
                b = Math.max(0, Math.min(255, b));

                result.setRGB(x, y, new Color(r, g, b).getRGB());
            }
        }
        return result;
    }

    // 降噪滤镜（使用高斯模糊）
    private static BufferedImage applyDenoiseFilter(BufferedImage source) {
        float[] matrix = {
                1 / 16f, 2 / 16f, 1 / 16f,
                2 / 16f, 4 / 16f, 2 / 16f,
                1 / 16f, 2 / 16f, 1 / 16f
        };

        BufferedImage result = new BufferedImage(
                source.getWidth(),
                source.getHeight(),
                BufferedImage.TYPE_INT_RGB);

        Kernel kernel = new Kernel(3, 3, matrix);
        ConvolveOp op = new ConvolveOp(kernel, ConvolveOp.EDGE_NO_OP, null);
        op.filter(source, result);

        return result;
    }

    // 色彩平衡滤镜
    private static BufferedImage applyColorBalanceFilter(BufferedImage source) {
        BufferedImage result = new BufferedImage(
                source.getWidth(),
                source.getHeight(),
                BufferedImage.TYPE_INT_RGB);

        long totalR = 0, totalG = 0, totalB = 0;
        int pixels = source.getWidth() * source.getHeight();

        for (int y = 0; y < source.getHeight(); y++) {
            for (int x = 0; x < source.getWidth(); x++) {
                Color color = new Color(source.getRGB(x, y));
                totalR += color.getRed();
                totalG += color.getGreen();
                totalB += color.getBlue();
            }
        }

        double avgR = (double) totalR / pixels;
        double avgG = (double) totalG / pixels;
        double avgB = (double) totalB / pixels;
        double avgTotal = (avgR + avgG + avgB) / 3;

        // 计算调整系数
        float rFactor = (float) (avgTotal / avgR);
        float gFactor = (float) (avgTotal / avgG);
        float bFactor = (float) (avgTotal / avgB);

        // 应用色彩平衡
        for (int y = 0; y < source.getHeight(); y++) {
            for (int x = 0; x < source.getWidth(); x++) {
                Color color = new Color(source.getRGB(x, y));

                int r = Math.min(255, (int) (color.getRed() * rFactor));
                int g = Math.min(255, (int) (color.getGreen() * gFactor));
                int b = Math.min(255, (int) (color.getBlue() * bFactor));

                result.setRGB(x, y, new Color(r, g, b).getRGB());
            }
        }
        return result;
    }
}